from typing import Optional, Generator

from langchain_core.messages import HumanMessage
from langgraph.types import StreamWriter
from pydantic import BaseModel
from langchain_core.runnables import RunnableConfig

from ...prompts.template import get_prompt_template
from ...utils.image import generate_image_file, get_image_prefix
from ...utils.str import get_message_content, print_message
from .base import BaseNodes, tool_node
from ...models.graph import GraphState, GraphRequest, AgentStepLog, GraphNodeParam
from ...llms import LLMContentModality,InsCodeModel
from loguru import logger




class ImageChatSchema(BaseModel):
    question: str


class ImageGenerateSchema(BaseModel):
    image_desc: str
    sizes: str
    type: str = "text-to-image"  # "text-to-image" or "image-to-image"


class ImageToPageSchema(BaseModel):
    question: str
    path: str


class ImageNodes(BaseNodes):

    def _get_vision_model(self)->InsCodeModel:
        """获取支持图像输入的模型，优先使用llm_chat"""
        if (self.llm_chat and LLMContentModality.IMAGE in self.llm_chat.input_modality):
            return self.llm_chat
        elif self.llm_vl:
            return self.llm_vl
        return None

    def _get_image_generation_model(self)->InsCodeModel:
        """获取支持图像生成的模型，优先使用llm_chat"""
        if (self.llm_chat and LLMContentModality.IMAGE in self.llm_chat.output_modality):
            return self.llm_chat
        elif self.llm_image:
            return self.llm_image
        return None

    @tool_node(args_schema=ImageChatSchema)
    def image_chat(self, state: GraphState, writer: StreamWriter, *, config: Optional[RunnableConfig] = None):
        """
        与图像对话。通过自然语言描述、分析或询问图像相关信息。
        参数:
        - question: (必填) 关于图像的具体问题
        用法:
        <image_chat>
        <question>这张图片中有什么物体？</question>
        </image_chat>
        """
        vision_model = self._get_vision_model()
        if not vision_model:
            logger.error("Not provided with a vision language model")
            return self.node_failure(message="Not provided with a vision language model")

        try:
            request = GraphRequest.from_runnable_config(config)
            question = state.goto.params.get("question", "describe the image")
            logger.info(f"Processing image_chat with question: '{question}' using model: {vision_model.model_name}")

            messages = [{
                "role": "user",
                "content": [
                    {"type": "text", "text": question},
                    {"type": "image_url", "image_url": {"url": f'{get_image_prefix(request.image)}{request.image}'}}
                ]
            }]

            last_line = ""
            total_usage = None
            for chunk in vision_model.chat_model.stream(messages):
                last_line += chunk.content
                writer(AgentStepLog.build_answer(chunk=last_line, meta={"finish": False}))

                if hasattr(chunk, 'usage_metadata') and chunk.usage_metadata:
                    total_usage = chunk.usage_metadata

            if total_usage:
                state.add_usage_meta(vision_model.model_name, total_usage)

            writer(AgentStepLog.build_answer(chunk=last_line, meta={"finish": True}))
            return self.node_success(message=last_line)
        except Exception as e:
            logger.error(f"Error in image chat: {e}")
            return self.node_failure(message="Error in image chat")

    @tool_node(args_schema=ImageGenerateSchema, stream=False)
    def image_generate(self, state: GraphState, writer: StreamWriter, *, config: Optional[RunnableConfig] = None):
        """
        AI图像生成。根据详细描述生成指定尺寸的高质量图像。
        参数:
        - image_desc: (必填) 图像详细描述(风格、主题、内容等)
        - sizes: (可选) 图像尺寸，格式"宽x高"，默认"1024x1024"。如需多个尺寸用逗号分隔，如"512x512,1024x1024"
        - type: (可选) 图像生成类型，"text-to-image"(文生图)或"image-to-image"(图生图)，默认为"text-to-image"
        用法:
        <image_generate>
        <image_desc>印象派风格的日落山景，雪山倒映在湖面，暖橙紫色天空</image_desc>
        <sizes>1024x1024</sizes>
        <type>text-to-image</type>
        </image_generate>
        """
        image_model = self._get_image_generation_model()
        if not image_model:
            logger.error("Not provided with an image generation model")
            return self.node_failure(message="Not provided with an image generation model")

        image_desc = state.goto.params.get("image_desc", "generate an blank image")
        sizes = state.goto.params.get("sizes", "1024x1024")
        generation_type = state.goto.params.get("type", "text-to-image")

        try:
            request = GraphRequest.from_runnable_config(config)
            image_size = [size.strip() for size in sizes.split(",")]
            logger.info(f"Generating image with description: '{image_desc}' using model: {image_model.model_name}, type: {generation_type}")

            result = generate_image_file(
                llm=image_model.chat_model,
                image_desc=image_desc,
                image_size=image_size,
                input_images=request.image if generation_type == "image-to-image" else None,
                valid_image_sizes=image_model.valid_image_sizes,
                image_type="png"
            )

            if isinstance(result, tuple) and len(result) == 2:
                images, usage_metadata = result
                if usage_metadata:
                    state.add_usage_meta(image_model.model_name, usage_metadata)
            else:
                images = result
                state.add_image_usage(image_model.model_name, len(images) if isinstance(images, list) else 1)

            logger.info(f"Image generation successful: {images}")
            return self.node_success(message=f"Image generated successfully. {images}")
        except Exception as e:
            logger.error(f"Error in image generation: {e}")
            return self.node_failure(message="Error in image generation")

    @tool_node(args_schema=ImageToPageSchema)
    def image_to_page(self, state: GraphState, writer: StreamWriter, *, config: Optional[RunnableConfig] = None):
        """
        图像转网页。基于上传图像自动生成对应的HTML页面实现。
        参数:
        - path: (必填) 生成页面的保存路径(相对当前目录)
        用法:
        <image_to_page>
        <path>generated_page.html</path>
        </image_to_page>
        """
        vision_model = self._get_vision_model()
        if not vision_model:
            logger.error("Not provided with a vision language model")
            return self.node_failure(message="Not provided with a vision language model")

        try:
            request = GraphRequest.from_runnable_config(config)
            path = state.goto.params.get("path", "generated_page.html")

            question = "Generate a page based on the image"
            for msg in state.messages:
                if (isinstance(msg, HumanMessage) and
                    msg.additional_kwargs and
                    msg.additional_kwargs.get("is_question", False)):
                    question = get_message_content(msg)
                    break

            if not request.image:
                logger.error("No image provided for image_to_page")
                return self.node_failure(message="No image provided")

            logger.info(f"Generating page from image using model: {vision_model.model_name}, question: {question}")
            action = "write_to_file"
            last_line = None
            total_usage = None

            for line, usage in ai_generate_page_by_image(question=question, image=request.image, llm=vision_model.chat_model, verbose=request.verbose):
                if not line:
                    continue
                last_line = line
                if usage:
                    total_usage = usage
                writer(AgentStepLog.build_tool_action(action=action,
                                                      output={"type": action,
                                                              "params": {"path": path, "content": last_line}},
                                                      meta={"finish": False}))

            if total_usage:
                state.add_usage_meta(vision_model.model_name, total_usage)

            writer(AgentStepLog.build_tool_action(action=action,
                                                  output={"type": action,
                                                          "params": {"path": path, "content": last_line}},
                                                  meta={"finish": True}))
            logger.info(f"Page generation completed successfully, saved to: {path}")
            return self.node_success(message=f"Page generated successfully at {path}, page content:\n\n {last_line}",
                                     force_goto=GraphNodeParam(type=action))
        except Exception as e:
            logger.error(f"Error in image to page: {e}")
            return self.node_failure(message="Error in image to page")


def ai_generate_page_by_image(question: str,
                              image: str,
                              llm,
                              verbose: bool = False,
                              **kwargs) -> Generator:
    """
    根据图片生成页面代码
    """

    def _extract_page_content(text: str) -> str:
        if "<page_content>" in text:
            start = text.find("<page_content>") + len("<page_content>")
            if "</page_content>" in text:
                end = text.find("</page_content>")
                return text[start:end].strip()
            return text[start:].strip()
        return ""

    try:
        messages = [{
            "role": "user",
            "content": [
                {"type": "text", "text": get_prompt_template("tool_image2page", question=question)},
                {"type": "image_url", "image_url": {"url": f'{get_image_prefix(image)}{image}'}}
            ]
        }]
        
        if verbose:
            print_message(messages)

        last_line = ""
        last_content = ""
        total_usage = None

        for chunk in llm.stream(messages):
            last_line += chunk.content
            # 收集usage信息
            if hasattr(chunk, 'usage_metadata') and chunk.usage_metadata:
                total_usage = chunk.usage_metadata

            if len(last_line) < 30:
                continue

            output = last_line[:-30]
            content = _extract_page_content(output)
            if content and content != last_content:
                yield content, total_usage
                last_content = content

        if last_line:
            content = _extract_page_content(last_line)
            if content and content != last_content:
                yield content, total_usage

        logger.info("Page generation from image completed successfully")
    except Exception as e:
        logger.error(f"Failed to ai_generate_page_by_image: {e}")
        raise Exception("Failed to generate page by image")
